from langchain.document_loaders import TextLoader, UnstructuredFileLoader
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter, CharacterTextSplitter
from langchain.vectorstores import Chroma
import sentence_transformers


# 导入文本
loader = UnstructuredFileLoader("压铸机2.txt", mode="elements")  # 使用 UnstructuredFileLoader 加载器以元素模式加载文件
text_splitor = CharacterTextSplitter()      # 使用 CharacterTextSplitter 来分割文件中的文本
docs = loader.load_and_split(text_splitor)  # 加载文件并进行文本分割

embedding_model_dict = {
    "text2vec": "C:\\Work\\llm\\text2vec-large-chinese"
}

EMBEDDING_MODEL = "text2vec"
# 初始化 hugginFace 的 embeddings 对象
embeddings = HuggingFaceEmbeddings(model_name=embedding_model_dict[EMBEDDING_MODEL])
embeddings.client = sentence_transformers.SentenceTransformer(
    embeddings.model_name, device='cuda')



# 初始化加载器
db = Chroma.from_documents(docs, embeddings, persist_directory="./chroma/casting2")
# # 持久化
db.persist()

question = "机床压力异常，不增压"
similarDocs = db.similarity_search(question, include_metadata=True,k=4)
[print(x) for x in similarDocs]
